"""
一个使用 LangChain 构建的聊天机器人，支持消息历史和流式输出
"""

import os
from typing import Dict, List, Optional
from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.llms.tongyi import Tongyi
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tracers import ConsoleCallbackHandler
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

# 加载环境变量
load_dotenv()

class ChatBot:
    def __init__(self, debug: bool = False):
        """
        初始化聊天机器人
        
        Args:
            debug: 是否启用调试模式，启用后会在控制台显示详细的调用信息
        """
        self.api_key = os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("请设置 DASHSCOPE_API_KEY 环境变量")
        
        # 设置回调处理器
        self.callbacks = [ConsoleCallbackHandler()] if debug else None
        
        # 初始化通义大模型
        self.llm = Tongyi(
            model_name="qwen-plus",  # 可选: qwen-turbo, qwen-plus, qwen-max
            temperature=0.7,
            api_key=self.api_key,
            callbacks=self.callbacks
        )

        # 创建提示词模板
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个乐于助人的 AI 助手。请用简洁明了的方式回答问题。"),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{input}")
        ])

        # 创建输出解析器
        self.output_parser = StrOutputParser()

        # 构建 LCEL 链
        self.chain = (
            self.prompt 
            | self.llm 
            | self.output_parser
        )

    def run(self, user_input: str, history: List[Dict] = None) -> str:
        """
        运行聊天机器人
        
        Args:
            user_input: 用户输入的问题
            history: 对话历史，格式为 [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
            
        Returns:
            str: AI 助手的回答
            
        Raises:
            Exception: 当处理过程中发生错误时
        """
        try:
            # 转换历史记录格式
            messages = []
            if history:
                for msg in history:
                    if msg["role"] == "user":
                        messages.append(HumanMessage(content=msg["content"]))
                    elif msg["role"] == "assistant":
                        messages.append(AIMessage(content=msg["content"]))
            
            response = self.chain.invoke(
                {
                    "input": user_input,
                    "history": messages
                },
                config={
                    "callbacks": self.callbacks,
                    "metadata": {"user_input": user_input}
                }
            )
            return response
        except Exception as e:
            return f"发生错误: {str(e)}"

    def get_chain(self):
        """获取 LCEL 链，用于 LangServe 部署"""
        return self.chain

def create_chatbot(debug: bool = False) -> ChatBot:
    """
    创建聊天机器人实例
    
    Args:
        debug: 是否启用调试模式
        
    Returns:
        ChatBot: 聊天机器人实例
    """
    return ChatBot(debug=debug)

def main():
    """主函数"""
    chatbot = create_chatbot(debug=True)
    
    print("欢迎使用通义大模型聊天机器人！输入 'quit' 退出。")
    print("当前使用的是 qwen-plus 模型，您可以在代码中修改为 qwen-turbo 或 qwen-max")
    
    history = []
    while True:
        user_input = input("\n请输入您的问题: ")
        if user_input.lower() == 'quit':
            break
            
        response = chatbot.run(user_input, history)
        print(f"\nAI 助手: {response}")
        
        # 更新历史记录
        history.append({"role": "user", "content": user_input})
        history.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main() 